# SPDX-License-Identifier: LGPL-3.0-or-later
import unittest

import torch

from deepmd.pt.cxx_op import (
    ENABLE_CUSTOMIZED_OP,
)
from deepmd.pt.utils import (
    env,
)

from ..consistent.common import (
    parameterized,
)


@parameterized((torch.float64, torch.float32))
@unittest.skipIf(not ENABLE_CUSTOMIZED_OP, "PyTorch customized OPs are not built")
class TestTabulateFusionSeTOp(unittest.TestCase):
    def setUp(self) -> None:
        (dtype,) = self.param
        if dtype == torch.float64:
            self.prec = 1e-10
        elif dtype == torch.float32:
            # JZ: not sure the reason, but 1e-5 cannot pass the grad test
            self.prec = 1e-3
        self.table_tensor = torch.tensor(
            [
                -1.0600000163027882e02,
                7.7059358807135015e02,
                -5.6954714749735385e03,
                1.2167808756610991e03,
                -7.6199102434332218e01,
                1.0706136029373441e00,
                -1.0600000164528124e02,
                7.7059358630452323e02,
                -5.6954715659539552e03,
                1.2167808757436076e03,
                -7.6199099707724926e01,
                1.0706134206080884e00,
                -1.0600000163027882e02,
                7.7059358807135015e02,
                -5.6954714749735385e03,
                1.2167808756610991e03,
                -7.6199102434332218e01,
                1.0706136029373441e00,
                -1.0600000164528124e02,
                7.7059358630452323e02,
                -5.6954715659539552e03,
                1.2167808757436076e03,
                -7.6199099707724926e01,
                1.0706134206080884e00,
                -9.6000006759336443e01,
                6.2969719646863621e02,
                -4.2053706363664551e03,
                9.0372155784831205e02,
                -5.7600014239472898e01,
                8.6528676197113796e-01,
                -9.6000006828502180e01,
                6.2969718981238339e02,
                -4.2053709121998018e03,
                9.0372156236848912e02,
                -5.7600006817493266e01,
                8.6528625106787871e-01,
                -9.6000006759336443e01,
                6.2969719646863621e02,
                -4.2053706363664551e03,
                9.0372155784831205e02,
                -5.7600014239472898e01,
                8.6528676197113796e-01,
                -9.6000006828502180e01,
                6.2969718981238339e02,
                -4.2053709121998018e03,
                9.0372156236848912e02,
                -5.7600006817493266e01,
                8.6528625106787871e-01,
                -8.6000028021606425e01,
                5.0303296429845562e02,
                -3.0008648248894533e03,
                6.4939597734382562e02,
                -4.2250984019314707e01,
                6.8180015607155764e-01,
                -8.6000028340480625e01,
                5.0303293978396903e02,
                -3.0008656209622986e03,
                6.4939600529391078e02,
                -4.2250965541906716e01,
                6.8179882734268982e-01,
                -8.6000028021606425e01,
                5.0303296429845562e02,
                -3.0008648248894533e03,
                6.4939597734382562e02,
                -4.2250984019314707e01,
                6.8180015607155764e-01,
                -8.6000028340480625e01,
                5.0303293978396903e02,
                -3.0008656209622986e03,
                6.4939600529353049e02,
                -4.2250965541830588e01,
                6.8179882733888086e-01,
                -7.6000116148038558e01,
                3.9060139597613619e02,
                -2.0515743554479322e03,
                4.4772754091167945e02,
                -2.9848087537832814e01,
                5.2014755686537917e-01,
                -7.6000117618125429e01,
                3.9060130821883052e02,
                -2.0515765138621105e03,
                4.4772766653712006e02,
                -2.9848047259266409e01,
                5.2014443989116910e-01,
                -7.6000116148038558e01,
                3.9060139597613619e02,
                -2.0515743554479322e03,
                4.4772754091167945e02,
                -2.9848087537832814e01,
                5.2014755686537917e-01,
                -7.6000117618125742e01,
                3.9060130821877993e02,
                -2.0515765138659344e03,
                4.4772766652483722e02,
                -2.9848047256692499e01,
                5.2014443976043645e-01,
                -6.6000481290731443e01,
                2.9240425245900917e02,
                -1.3271250821434478e03,
                2.9263955624337893e02,
                -2.0087224005740719e01,
                3.8031147992206349e-01,
                -6.6000488067863742e01,
                2.9240394960550276e02,
                -1.3271304743966571e03,
                2.9264002765325057e02,
                -2.0087154325946980e01,
                3.8030522013794582e-01,
                -6.6000481290731443e01,
                2.9240425245900917e02,
                -1.3271250821434478e03,
                2.9263955624337893e02,
                -2.0087224005740719e01,
                3.8031147992206349e-01,
                -6.6000488067883694e01,
                2.9240394960308691e02,
                -1.3271304745319526e03,
                2.9264002727267626e02,
                -2.0087154245656002e01,
                3.8030521605011575e-01,
                -5.6001992867343972e01,
                2.0844745574402617e02,
                -7.9715799906587699e02,
                1.7805563184427194e02,
                -1.2663929104029080e01,
                2.6224978307822894e-01,
                -5.6002024103130161e01,
                2.0844646075692629e02,
                -7.9717003898786652e02,
                1.7805715054974732e02,
                -1.2663864677938077e01,
                2.6224029170957303e-01,
                -5.6001992867343972e01,
                2.0844745574402617e02,
                -7.9715799906587699e02,
                1.7805563184427194e02,
                -1.2663929104029080e01,
                2.6224978307822894e-01,
                -5.6002024104383771e01,
                2.0844646064871867e02,
                -7.9717004324410516e02,
                1.7805714044473001e02,
                -1.2663862524337585e01,
                2.6224018166598279e-01,
                -4.6008230210744550e01,
                1.3874976550319553e02,
                -4.3134867537287749e02,
                9.7902623595157010e01,
                -7.2734403121911884e00,
                1.6589123996688057e-01,
                -4.6008373996710617e01,
                1.3874671965012058e02,
                -4.3137141216256458e02,
                9.7906861443792735e01,
                -7.2735856084076280e00,
                1.6588642735924275e-01,
                -4.6008230210744550e01,
                1.3874976550319553e02,
                -4.3134867537287749e02,
                9.7902623595157010e01,
                -7.2734403121911884e00,
                1.6589123996688057e-01,
                -4.6008374075307870e01,
                1.3874671513440606e02,
                -4.3137152784492957e02,
                9.7906652364871050e01,
                -7.2735401377994249e00,
                1.6588408717348646e-01,
                -3.6033642533368131e01,
                8.3364086172019398e01,
                -1.9942175516407502e02,
                4.6124022747838069e01,
                -3.6130563858549958e00,
                9.1249773312287188e-02,
                -3.6034298111245583e01,
                8.3355843868269616e01,
                -1.9945266030093268e02,
                4.6135000705962462e01,
                -3.6142786797647353e00,
                9.1293932043118198e-02,
                -3.6033642533368131e01,
                8.3364086172019398e01,
                -1.9942175516407502e02,
                4.6124022747838069e01,
                -3.6130563858549958e00,
                9.1249773312287188e-02,
                -3.6034302998781108e01,
                8.3355675173745269e01,
                -1.9945516784358935e02,
                4.6132303200740992e01,
                -3.6136582565667807e00,
                9.1261386291659793e-02,
                -2.6132076703837274e01,
                4.2398929436319683e01,
                -7.1037171119057973e01,
                1.3425662262407457e01,
                -7.5172495708992593e-01,
                7.7522572203268742e-03,
                -2.6134776894873077e01,
                4.2384732735328775e01,
                -7.1030526549717337e01,
                1.3431455085299461e01,
                -7.5302028721199155e-01,
                7.8186246126207160e-03,
                -2.6132076703837274e01,
                4.2398929436319683e01,
                -7.1037171119057973e01,
                1.3425662262405055e01,
                -7.5172495708944420e-01,
                7.7522572203027138e-03,
                -2.6135071381093578e01,
                4.2379566840123424e01,
                -7.1067162844830236e01,
                1.3434603316099608e01,
                -7.5251233833488806e-01,
                7.7734884077347950e-03,
                -2.2221480705551805e01,
                3.0067218434037404e01,
                -4.1779705297521097e01,
                -1.9077757705724110e02,
                3.6413466026808294e02,
                -1.6067397401486718e02,
                -2.2225430071703467e01,
                3.0060809113889512e01,
                -4.1712800191721314e01,
                -1.9084786311022177e02,
                3.6410062714257685e02,
                -1.6063028238785057e02,
                -2.2221480705551830e01,
                3.0067218434036263e01,
                -4.1779705297545611e01,
                -1.9077757705723738e02,
                3.6413466026815809e02,
                -1.6067397401492047e02,
                -2.2226913938674084e01,
                3.0042371820589185e01,
                -4.1801582285426832e01,
                -1.9048619249019526e02,
                3.6373874557858261e02,
                -1.6052358406417352e02,
                -2.1250858373060836e01,
                2.7343847665267702e01,
                -3.6044215009418814e01,
                -1.7618484800469861e02,
                3.3120085405644409e02,
                -1.4534825256321494e02,
                -2.1254939505030809e01,
                2.7342716030835884e01,
                -3.5955450545431681e01,
                -1.7635550119316844e02,
                3.3127447930769307e02,
                -1.4533876561022046e02,
                -2.1250858373060954e01,
                2.7343847665262818e01,
                -3.6044215009514119e01,
                -1.7618484800464822e02,
                3.3120085405666612e02,
                -1.4534825256338749e02,
                -2.1257155379297881e01,
                2.7317691772612619e01,
                -3.6063526926252166e01,
                -1.7588696592837897e02,
                3.3079005662384850e02,
                -1.4519086534447842e02,
                -2.0283472228681301e01,
                2.4763027042036295e01,
                -3.0876160316998963e01,
                -1.6184864900381874e02,
                2.9976970905591691e02,
                -1.3084395423768876e02,
                -2.0287461515322455e01,
                2.4769400540137131e01,
                -3.0762734380983186e01,
                -1.6214886052089241e02,
                2.9998995088792128e02,
                -1.3088331758129965e02,
                -2.0283472228681809e01,
                2.4763027042017129e01,
                -3.0876160317336627e01,
                -1.6184864900359682e02,
                2.9976970905662938e02,
                -1.3084395423826805e02,
                -2.0290765181946348e01,
                2.4735639907973120e01,
                -3.0892738413082597e01,
                -1.6154574482310053e02,
                2.9934595420013272e02,
                -1.3068028494926122e02,
                -1.9319499689234629e01,
                2.2323824431805683e01,
                -2.6243395369841849e01,
                -1.4782286378121026e02,
                2.6985759662396487e02,
                -1.1715474197881395e02,
                -1.9323022570439292e01,
                2.2340565860680357e01,
                -2.6102786429129356e01,
                -1.4828764857305418e02,
                2.7027298759214750e02,
                -1.1726163007473576e02,
                -1.9319499689236839e01,
                2.2323824431730525e01,
                -2.6243395371031539e01,
                -1.4782286378021576e02,
                2.6985759662609979e02,
                -1.1715474198068593e02,
                -1.9327939259284843e01,
                2.2295320666731183e01,
                -2.6257097174199931e01,
                -1.4751677383623073e02,
                2.6942341041084092e02,
                -1.1698575776762208e02,
                -1.8359079763330211e01,
                2.0025118950280675e01,
                -2.2113826757823226e01,
                -1.3415932552431914e02,
                2.4147795894487624e02,
                -1.0427314537549884e02,
                -1.8361534194530734e01,
                2.0055847278170305e01,
                -2.1944107342764479e01,
                -1.3482982214648752e02,
                2.4214772485703989e02,
                -1.0447085300268679e02,
                -1.8359079763339750e01,
                2.0025118949989704e01,
                -2.2113826761939308e01,
                -1.3415932552009582e02,
                2.4147795895089951e02,
                -1.0427314538136979e02,
                -1.8368836959765495e01,
                1.9995657614892380e01,
                -2.2124533894067383e01,
                -1.3385233293246981e02,
                2.4103659293914149e02,
                -1.0410011400771683e02,
                -1.7402299525814517e01,
                1.7865597763687486e01,
                -1.8455503416511757e01,
                -1.2090765118569301e02,
                2.1464125749038132e02,
                -9.2190581022134992e01,
                -1.7402744551259310e01,
                1.7914800567904472e01,
                -1.8255754666855470e01,
                -1.2183089355280822e02,
                2.1563582256173194e02,
                -9.2507405324257306e01,
                -1.7402299525855486e01,
                1.7865597762572605e01,
                -1.8455503430527756e01,
                -1.2090765116826699e02,
                2.1464125750558804e02,
                -9.2190581039770791e01,
                -1.7413567239985614e01,
                1.7835392747330133e01,
                -1.8463115133795956e01,
                -1.2060260469703572e02,
                2.1419685510959093e02,
                -9.2015134441585104e01,
                -1.6449179896085464e01,
                1.5843762224435309e01,
                -1.5236722252652665e01,
                -1.0811515163854509e02,
                1.8935506712501905e02,
                -8.0897437157402223e01,
                -1.6446174965543889e01,
                1.5916874201410112e01,
                -1.5007553197461570e01,
                -1.0934291295595986e02,
                1.9075532567542470e02,
                -8.1366596347119696e01,
                -1.6449179896260411e01,
                1.5843762220214204e01,
                -1.5236722299508587e01,
                -1.0811515156878269e02,
                1.8935506715588940e02,
                -8.0897437207525684e01,
                -1.6462173655481337e01,
                1.5813096619069219e01,
                -1.5241142983208677e01,
                -1.0781563484017332e02,
                1.8891289499393798e02,
                -8.0721658713418606e01,
                -1.5499661595231082e01,
                1.3957945516559789e01,
                -1.2426145992195885e01,
                -9.5826844741964834e01,
                1.6562434781973772e02,
                -7.0383233416004117e01,
                -1.5491037589250178e01,
                1.4061349904707843e01,
                -1.2170301483989650e01,
                -9.7412966929875139e01,
                1.6751874597575440e02,
                -7.1041920384880939e01,
                -1.5499661595973759e01,
                1.3957945500778198e01,
                -1.2426146145776961e01,
                -9.5826844470313858e01,
                1.6562434784656404e02,
                -7.0383233547510557e01,
                -1.5514618579274794e01,
                1.3927192540790591e01,
                -1.2427264674287118e01,
                -9.5537423121432880e01,
                1.6519113036542510e02,
                -7.0209783384625098e01,
                -1.4553592409098401e01,
                1.2206343505203831e01,
                -9.9929274597052196e00,
                -8.4085595900823435e01,
                1.4345191724964303e02,
                -6.0636862050381758e01,
                -1.4536130507533649e01,
                1.2347228125716077e01,
                -9.7159302678980044e00,
                -8.6081002959763751e01,
                1.4592996741513730e02,
                -6.1523840242331410e01,
                -1.4553592412232879e01,
                1.2206343446986155e01,
                -9.9929279524397305e00,
                -8.4085594870780753e01,
                1.4345191706222485e02,
                -6.0636862352071532e01,
                -1.4570766853404239e01,
                1.2175998366492486e01,
                -9.9905856922863112e00,
                -8.3812185051328299e01,
                1.4303633648493073e02,
                -6.0469165577726159e01,
                -1.3610717065161962e01,
                1.0587059629986399e01,
                -7.9068321681349163e00,
                -7.2932404423885004e01,
                1.2283913327111270e02,
                -5.1646910322317169e01,
                -1.3579708436673444e01,
                1.0773027159520954e01,
                -7.6175370796795425e00,
                -7.5376833196183071e01,
                1.2597958225245242e02,
                -5.2797863799745748e01,
                -1.3610717078313911e01,
                1.0587059418306087e01,
                -7.9068337121483454e00,
                -7.2932400620636059e01,
                1.2283913169238102e02,
                -5.1646910832841897e01,
                -1.3630368323321786e01,
                1.0557789879027116e01,
                -7.9007777139483810e00,
                -7.2682825476758552e01,
                1.2245259140017740e02,
                -5.1489446559796768e01,
                -1.2670671078399982e01,
                9.0981634949263963e00,
                -6.1383490362855788e00,
                -6.2406844162279825e01,
                1.0378677653422224e02,
                -4.3402055519687693e01,
                -1.2619333100308433e01,
                9.3364634226935799e00,
                -5.8491811509717584e00,
                -6.5316414528433455e01,
                1.0763857666200300e02,
                -4.4841832720191050e01,
                -1.2670671133253135e01,
                9.0981627374157021e00,
                -6.1383537481895356e00,
                -6.2406830503476570e01,
                1.0378676818216074e02,
                -4.3402055529436716e01,
                -1.2693036794620980e01,
                9.0708908225804148e00,
                -6.1281713411274001e00,
                -6.2191660620037396e01,
                1.0344456594081470e02,
                -4.3260806640248063e01,
                -1.1732979767504439e01,
                7.7377614739662697e00,
                -4.6587775146685351e00,
                -5.2547655563671029e01,
                8.6296103981829802e01,
                -3.5891515805495345e01,
                -1.1651721415208119e01,
                8.0340005825064456e00,
                -4.3852919661646119e00,
                -5.5898160750405737e01,
                9.0851291378134590e01,
                -3.7622755083739385e01,
                -1.1732979994779518e01,
                7.7377588120662892e00,
                -4.6587914600219875e00,
                -5.2547607987974565e01,
                8.6296066930227624e01,
                -3.5891510429190419e01,
                -1.1758218632638741e01,
                7.7137968422318544e00,
                -4.6438239588320966e00,
                -5.2381405657406454e01,
                8.6019170302439520e01,
                -3.5774653697918737e01,
                -1.0797063195543267e01,
                6.5040766534586290e00,
                -3.4402783696562169e00,
                -4.3393478931462226e01,
                7.0370032342568010e01,
                -2.9105535302381853e01,
                -1.0672637254876815e01,
                6.8603244928014488e00,
                -3.1995767859681346e00,
                -4.7101348454718874e01,
                7.5530774605740319e01,
                -3.1094453979913311e01,
                -1.0797064129672576e01,
                6.5040675030570139e00,
                -3.4403181344841500e00,
                -4.3393319126804485e01,
                7.0369884883020177e01,
                -2.9105501594155889e01,
                -1.0825134802124644e01,
                6.4853446725127366e00,
                -3.4195560956016346e00,
                -4.3296381389022351e01,
                7.0187483762520671e01,
                -2.9024415860031247e01,
                -9.8622468030169337e00,
                5.3955359781222549e00,
                -2.4558741324534137e00,
                -3.4983728078555984e01,
                5.6014425934291204e01,
                -2.3035887876475471e01,
                -9.6769173769353625e00,
                5.8079540801032961e00,
                -2.2635143148159220e00,
                -3.8890523502249145e01,
                6.1563046720547966e01,
                -2.5198820521877391e01,
                -9.8622505990399034e00,
                5.3955054149765509e00,
                -2.4559821583353774e00,
                -3.4983216045684472e01,
                5.6013889382190079e01,
                -2.3035736114340502e01,
                -9.8926597117464805e00,
                5.3849440641688187e00,
                -2.4279562878572039e00,
                -3.4983707025980287e01,
                5.5966629574570753e01,
                -2.3006306589550750e01,
                -8.9277749780883457e00,
                4.4108678323349286e00,
                -1.6793815271288624e00,
                -2.7359655656676122e01,
                4.3239544183593061e01,
                -1.7676416286664047e01,
                -8.6587749152265552e00,
                4.8674392165289442e00,
                -1.5450097170494306e00,
                -3.1230915545542118e01,
                4.8829474992442343e01,
                -1.9874755288141955e01,
                -8.9277901202336185e00,
                4.4107699183102085e00,
                -1.6796551456533098e00,
                -2.7358123514289456e01,
                4.3237769027728554e01,
                -1.7675844947587926e01,
                -8.9590559763951383e00,
                4.4128957610428623e00,
                -1.6423658138809611e00,
                -2.7493743583145054e01,
                4.3380518846300511e01,
                -1.7719639183506050e01,
                -7.9928164326293913e00,
                3.5492331091008302e00,
                -1.0852462622393610e00,
                -2.0565792757352423e01,
                3.2061909496398073e01,
                -1.3023704651715642e01,
                -7.6125412569887647e00,
                4.0287966748633526e00,
                -1.0084592804412351e00,
                -2.4116992333062022e01,
                3.7252797603904497e01,
                -1.5077495076198684e01,
                -7.9928747817255603e00,
                3.5489404571097585e00,
                -1.0858609980296849e00,
                -2.0561701094768868e01,
                3.2056747083970720e01,
                -1.3021877019728107e01,
                -8.0213899495838241e00,
                3.5708128515175943e00,
                -1.0368753205735253e00,
                -2.0877831538201836e01,
                3.2456559535389509e01,
                -1.3165540198118645e01,
                -7.0564174984379102e00,
                2.8104770395789380e00,
                -6.4821407306458223e-01,
                -1.4652118176169953e01,
                2.2507145963021038e01,
                -9.0780963613608154e00,
                -6.5338936679228468e00,
                3.2846161494194233e00,
                -6.1760141818709846e-01,
                -1.7606122820367215e01,
                2.6855555289500277e01,
                -1.0803821410528570e01,
                -7.0566263531717324e00,
                2.8097184139861691e00,
                -6.4925197579297411e-01,
                -1.4643483271177150e01,
                2.2495243692983838e01,
                -9.0734373052814821e00,
                -7.0742646195707266e00,
                2.8621047467298468e00,
                -5.8641470402843421e-01,
                -1.5178915176777426e01,
                2.3211717123277591e01,
                -9.3414295847965061e00,
                -6.1172231064332783e00,
                2.1957964102200167e00,
                -3.4265643705632465e-01,
                -9.6769153352706798e00,
                1.4613873405033004e01,
                -5.8450824172251430e00,
                -5.4212678780860326e00,
                2.6341589573018260e00,
                -3.4085224757280796e-01,
                -1.1835854891340576e01,
                1.7794701474942944e01,
                -7.1075278532253687e00,
                -6.1178367984533244e00,
                2.1945528943967396e00,
                -3.4261268423617658e-01,
                -9.6695829134679272e00,
                1.4600877298870854e01,
                -5.8381668136523013e00,
                -6.1072022151656586e00,
                2.2922503774685161e00,
                -2.6715334266026142e-01,
                -1.0408120531614587e01,
                1.5617405440391840e01,
                -6.2270636615178061e00,
                -5.1722074807324017e00,
                1.7098190643016411e00,
                -1.4098618492175408e-01,
                -5.7061337346696464e00,
                8.4331806866534098e00,
                -3.3349192888568142e00,
                -4.2766424379800121e00,
                2.0860564217794284e00,
                -1.5548660419053545e-01,
                -7.0034949575065015e00,
                1.0332245608764421e01,
                -4.0873492185766374e00,
                -5.1727690165421372e00,
                1.7132539127425084e00,
                -1.2776576793785877e-01,
                -5.7565343018918274e00,
                8.4941254548170697e00,
                -3.3479852132230872e00,
                -5.0998839330979591e00,
                1.8678855512825561e00,
                -5.7718910331047868e-02,
                -6.5095346397755423e00,
                9.5462002113817768e00,
                -3.7632628689263172e00,
                -4.2112469382255613e00,
                1.3675717927787789e00,
                -9.4961575783498800e-03,
                -2.7877417589321136e00,
                3.9953503912711956e00,
                -1.5499906707437840e00,
                -3.1046711877098376e00,
                1.6568346830533449e00,
                -4.5990009889900242e-02,
                -3.3140676307068091e00,
                4.7472200808709299e00,
                -1.8492173878772247e00,
                -4.1976749320353317e00,
                1.4246952243441517e00,
                8.7531923058200650e-02,
                -3.0996975434049761e00,
                4.4668738099197531e00,
                -1.7103055321708385e00,
                -4.0163145894665320e00,
                1.5923303121893606e00,
                5.8249749369824022e-02,
                -3.3748048713195491e00,
                4.7925769874900315e00,
                -1.8598420111853879e00,
                -3.1955533414298376e00,
                1.2168024121915868e00,
                9.9474205814620603e-02,
                -8.6811124876189694e-01,
                1.1994338853723501e00,
                -4.4837238870567747e-01,
                -1.9098914522594992e00,
                1.3654451552507061e00,
                2.9537044429980407e-03,
                -9.3701125207094127e-01,
                1.2575365835116745e00,
                -4.7248060681970733e-01,
                -3.0285770502890443e00,
                1.6166340190704305e00,
                4.8662683065338386e-01,
                -1.2308607057515726e00,
                1.6114560066217587e00,
                -6.5896729332189652e-01,
                -2.8078044229222514e00,
                1.4555130910035559e00,
                9.0876948497501955e-02,
                -1.0566809618626720e00,
                1.3938154223720176e00,
                -5.2279617091852160e-01,
                -1.9963264755188566e00,
                1.3672906754961440e00,
                2.0801988470625002e-01,
                2.0083818728351077e-02,
                -1.5135587406137185e-02,
                -1.4175240342178652e-02,
                -6.9344786794476854e-01,
                1.2280621078720415e00,
                1.2333381103148277e-02,
                -1.0895386066093759e-02,
                2.1764282171790141e-02,
                -1.0106900291744604e-02,
                -1.2036881930169383e00,
                2.0482931230000392e00,
                -1.2689218008973949e-01,
                -5.0580690719339239e-01,
                3.4047786101030464e-01,
                -7.0959386937004015e-02,
                -1.4470760938303664e00,
                1.4285049373060201e00,
                5.5764887956399375e-02,
                -2.9461990750009881e-02,
                2.3005167601875431e-02,
                -1.0760396189439407e-02,
                -4.3024292433642597e-01,
                1.7121633497582587e00,
                3.5705413032693957e-02,
                -9.9216800479772127e-01,
                1.5115432403429119e00,
                -6.3985596276149748e-01,
                5.4770961684437192e-01,
                1.2565653391084903e00,
                9.1639130181564755e-03,
                -6.8547618650262643e-01,
                1.2037212931265591e00,
                -5.1526772142324506e-01,
                4.8142431677326969e-01,
                1.2842025505965851e00,
                -3.1103960497811806e-01,
                -3.8667287940463613e-01,
                9.2663039525338942e-01,
                -4.1330437951972537e-01,
                1.9976512094478704e-02,
                1.4898674304290889e00,
                -2.1940405767858565e-03,
                -8.0791207141984167e-01,
                1.3979310081478775e00,
                -5.9845265079421794e-01,
                1.1971451112382212e00,
                1.6539633089946477e00,
                -2.7009878691796618e-01,
                -2.8868139196850624e00,
                4.7294193613612734e00,
                -1.9578020397520424e00,
                1.8164162541717044e00,
                1.4570111710269262e00,
                2.2385898037164991e-02,
                -3.1195681762439769e00,
                4.9723722392038878e00,
                -2.0423972644796100e00,
                1.5812403987207633e00,
                1.1421043858413655e00,
                -4.4319666868952730e-02,
                -2.3144705949527720e00,
                3.7448930479898297e00,
                -1.5426803544433196e00,
                1.4992161878806018e00,
                1.6612039136364238e00,
                -2.2870713891204597e-02,
                -3.4442115437939465e00,
                5.5057190995408973e00,
                -2.2657208348376137e00,
                2.4658130352390710e00,
                1.5819912227884063e00,
                -1.3204477532594588e-01,
                -5.7752803465671017e00,
                9.0677018990478242e00,
                -3.6843468204828174e00,
                3.1062201217160963e00,
                1.8205810727868250e00,
                7.3942159732456811e-02,
                -7.3418038323250947e00,
                1.1309154676354810e01,
                -4.5733470083866452e00,
                2.5667672162869133e00,
                1.3762236869878626e00,
                5.4823291778512563e-02,
                -5.5558964069977943e00,
                8.5620133672289516e00,
                -3.4575259608624478e00,
                2.9333361085351610e00,
                1.9771000784477066e00,
                2.1600903596218385e-02,
                -7.7786452012965430e00,
                1.2026327126407146e01,
                -4.8722408979121159e00,
                3.5238342146994350e00,
                1.8411341262124141e00,
                1.0485737443151430e-01,
                -1.0316470080846322e01,
                1.5628354265192609e01,
                -6.2547428286449396e00,
                4.3947471898784478e00,
                2.3129375587624681e00,
                1.6998863701958250e-01,
                -1.3069120913924280e01,
                1.9764673064124775e01,
                -7.9234176878170990e00,
                3.5464051944219954e00,
                1.7786047141550632e00,
                1.8395466553434961e-01,
                -1.0256713338978345e01,
                1.5450540198835597e01,
                -6.1709943751208902e00,
                4.3074781177775723e00,
                2.4284702978185178e00,
                1.2121907902830774e-01,
                -1.3510697720561426e01,
                2.0490823414440431e01,
                -8.2265504110307699e00,
                4.5269670710447079e00,
                2.3411415500822019e00,
                3.7814443659878427e-01,
                -1.6533454371385766e01,
                2.4532574055181296e01,
                -9.7222898630871342e00,
                5.6498078480438974e00,
                2.8871559084424092e00,
                3.1648740182441881e-01,
                -1.9832336139347099e01,
                2.9630584562783888e01,
                -1.1804975183138390e01,
                4.5317970588477650e00,
                2.3235629480266455e00,
                4.0711209040396701e-01,
                -1.6523611973754900e01,
                2.4482080409856291e01,
                -9.6968326211377835e00,
                5.6107427774726322e00,
                2.9693568967987254e00,
                2.6856229367890733e-01,
                -2.0186235796983127e01,
                3.0228033555488111e01,
                -1.2057362656117963e01,
                5.5230828784340904e00,
                3.0159142144119913e00,
                7.5032702265793638e-01,
                -2.4452361306480910e01,
                3.5745746299744695e01,
                -1.4059387633540990e01,
                6.8467243986091164e00,
                3.5205846294935204e00,
                5.5323452910250115e-01,
                -2.7424447720726722e01,
                4.0542113968978946e01,
                -1.6058340606199877e01,
                5.5241079122419858e00,
                3.0111097413061287e00,
                7.6043241689918206e-01,
                -2.4453330947201032e01,
                3.5733842835424838e01,
                -1.4052622761934279e01,
                6.8330970703372866e00,
                3.5730950345697865e00,
                5.0442967447855436e-01,
                -2.7630302835415993e01,
                4.0921397061842079e01,
                -1.6223699529825666e01,
                6.5233214752268127e00,
                3.8455313715589599e00,
                1.2738445662734672e00,
                -3.4142511056048967e01,
                4.9288751118195229e01,
                -1.9258816488331760e01,
                7.9798691992574877e00,
                4.2304633704347614e00,
                9.4916911879724064e-01,
                -3.6082800915305256e01,
                5.2740474636382487e01,
                -2.0757970588732530e01,
                6.5235391967368317e00,
                3.8442392655293900e00,
                1.2772689685023881e00,
                -3.4144245582802192e01,
                4.9286600694030149e01,
                -1.9257235266278844e01,
                7.9780164759860508e00,
                4.2581364755189171e00,
                9.0490824102641643e-01,
                -3.6146890048111374e01,
                5.2902251888236343e01,
                -2.0834714063750525e01,
                7.5301209868737518e00,
                4.8266093670811516e00,
                1.9906532239804082e00,
                -4.5696171225139402e01,
                6.5222794336738914e01,
                -2.5330008845677121e01,
                9.0592048208341964e00,
                5.0524444639807982e00,
                1.5639083038511417e00,
                -4.6227354827270197e01,
                6.6742768625790532e01,
                -2.6090733281390481e01,
                7.5301672757177256e00,
                4.8262668988539703e00,
                1.9917837214882572e00,
                -4.5697152262800707e01,
                6.5222641787790508e01,
                -2.5329699752317662e01,
                9.0617089689058279e00,
                5.0627200474303731e00,
                1.5306087886050987e00,
                -4.6201245261995687e01,
                6.6753711704174307e01,
                -2.6103836713323240e01,
                8.5439978438576958e00,
                5.9605352581937785e00,
                2.9388171122244109e00,
                -5.9213652478598007e01,
                8.3623964589400401e01,
                -3.2288651007290504e01,
                1.0100238105795977e01,
                6.0156046860821641e00,
                2.4311227628788585e00,
                -5.8189717323516248e01,
                8.2972590004142106e01,
                -3.2212869674305303e01,
                8.5440076687321067e00,
                5.9604459430021439e00,
                2.9391801366526531e00,
                -5.9214078468041464e01,
                8.3624068891376510e01,
                -3.2288610777657510e01,
                1.0103667533796683e01,
                6.0158650887345448e00,
                2.4107760944314816e00,
                -5.8125625048064265e01,
                8.2906979417176174e01,
                -3.2191629006406409e01,
                9.5650113177877785e00,
                7.2498153679976820e00,
                4.1551371399277919e00,
                -7.4795843598083408e01,
                1.0457037732454131e02,
                -4.0151433068943419e01,
                1.1116968561077568e01,
                7.1347098863330896e00,
                3.5688140741297674e00,
                -7.2151486218593305e01,
                1.0165680693075836e02,
                -3.9206269356622016e01,
                9.5650133940644455e00,
                7.2497924894015711e00,
                4.1552503042122613e00,
                -7.4796005009548836e01,
                1.0457044971811401e02,
                -4.0151435976986221e01,
                1.1120034079668221e01,
                7.1303147700774092e00,
                3.5594873892317103e00,
                -7.2082067018068685e01,
                1.0156598726189708e02,
                -3.9171834664292227e01,
                1.0593064483227742e01,
                8.6969028070512202e00,
                5.6755396034912966e00,
                -9.2539537763180832e01,
                1.2813560149579646e02,
                -4.8933613418447223e01,
                1.2119543877083460e01,
                8.4137603187360543e00,
                4.9925034366798311e00,
                -8.8194505075704640e01,
                1.2287993196505218e02,
                -4.7096724506223822e01,
                1.0593064919257221e01,
                8.6968970567044934e00,
                5.6755738143875760e00,
                -9.2539593640863643e01,
                1.2813563331215474e02,
                -4.8933618162805772e01,
                1.2121921818513506e01,
                8.4078642204619420e00,
                4.9908632634858190e00,
                -8.8134432374832016e01,
                1.2279086550380391e02,
                -4.7060844505587738e01,
                1.1627957207938659e01,
                1.0303707615441018e01,
                7.5344011042552923e00,
                -1.1253294830348190e02,
                1.5438372244089408e02,
                -5.8647453529357783e01,
                1.3114510015623049e01,
                9.8513572940713416e00,
                6.7213349376406626e00,
                -1.0635738219113546e02,
                1.4665751311861146e02,
                -5.5881528760137869e01,
                1.1627957298834614e01,
                1.0303706197478814e01,
                7.5344111366673712e00,
                -1.1253296638384563e02,
                1.5438373415898508e02,
                -5.8647455853629580e01,
                1.3116237925845430e01,
                9.8455331102145145e00,
                6.7243141059359051e00,
                -1.0631074264006560e02,
                1.4658112805680690e02,
                -5.5849452095162235e01,
                1.2669386535689361e01,
                1.2071287030293307e01,
                9.7633555455962835e00,
                -1.3485075345900265e02,
                1.8336444946299886e02,
                -6.9300787627414508e01,
                1.4105804414673191e01,
                1.1444289269702800e01,
                8.7789794745243590e00,
                -1.2666835962860844e02,
                1.7298274034188972e02,
                -6.5547771558832267e01,
                1.2669386554490638e01,
                1.2071286687068984e01,
                9.7633584027450482e00,
                -1.3485075900242089e02,
                1.8336445335820781e02,
                -6.9300788508071975e01,
                1.4107018463574896e01,
                1.1439185153305873e01,
                8.7843335749580440e00,
                -1.2663444344319166e02,
                1.7292158897636148e02,
                -6.5521162694327174e01,
                1.3716937488160630e01,
                1.3999597459400730e01,
                1.2389915672436279e01,
                -1.5954894249539399e02,
                2.1510813446746886e02,
                -8.0895567204040049e01,
                1.5095682313349364e01,
                1.3189272906323732e01,
                1.1192627051714643e01,
                -1.4915916817312757e02,
                2.0184825850919157e02,
                -7.6081293415969839e01,
                1.3716937492019641e01,
                1.3999597377767842e01,
                1.2389916464009524e01,
                -1.5954894412085929e02,
                2.1510813567394996e02,
                -8.0895567498068928e01,
                1.5096520030681436e01,
                1.3185064407456906e01,
                1.1198910160279951e01,
                -1.4913565617175487e02,
                2.0180124290250004e02,
                -7.6060129778156622e01,
                1.4770075388032444e01,
                1.6087303167766446e01,
                1.5436222950666867e01,
                -1.8666021493779203e02,
                2.4962122089688103e02,
                -9.3426463524457304e01,
                1.6085379191481852e01,
                1.5083589447287226e01,
                1.3991739427782750e01,
                -1.7386892459375579e02,
                2.3325385095807121e02,
                -8.7470099643500802e01,
                1.4770075388818769e01,
                1.6087303148664304e01,
                1.5436223164442264e01,
                -1.8666021539675981e02,
                2.4962122125116741e02,
                -9.3426463615076329e01,
                1.6085951551006787e01,
                1.5080238931969067e01,
                1.3998101278449143e01,
                -1.7385331837944693e02,
                2.3321864790104019e02,
                -8.7453697552144448e01,
                1.5828143941097450e01,
                1.8331670220961666e01,
                1.8918268274003861e01,
                -2.1619095210442941e02,
                2.8688297635978756e02,
                -1.0687973526499771e02,
                1.7075534787366465e01,
                1.7125200136366264e01,
                1.7207074959934751e01,
                -2.0084388544719391e02,
                2.6720765911058965e02,
                -9.9705133726570395e01,
                1.5828143941256627e01,
                1.8331670216557445e01,
                1.8918268330404022e01,
                -2.1619095222989833e02,
                2.8688297645950814e02,
                -1.0687973529137253e02,
                1.7075923730873765e01,
                1.7122590193964911e01,
                1.7213058024904747e01,
                -2.0083402645820061e02,
                2.6718180837697332e02,
                -9.9692640534772679e01,
                1.6890371426423382e01,
                2.0728579569842751e01,
                2.2845917469463828e01,
                -2.4812083435502871e02,
                3.2684448823688496e02,
                -1.2123263616047282e02,
                1.8066449820492846e01,
                1.9312661524160735e01,
                2.0870036016187061e01,
                -2.3013589616073858e02,
                3.0372498377642154e02,
                -1.1277999824352135e02,
                1.6890371426455424e01,
                2.0728579568840633e01,
                2.2845917484032956e01,
                -2.4812083438838550e02,
                3.2684448826399682e02,
                -1.2123263616782057e02,
                1.8066713333743454e01,
                1.9310657703202459e01,
                2.0875423564416035e01,
                -2.3013008228413184e02,
                3.0370630494679148e02,
                -1.1277060230387309e02,
                1.7955886187113396e01,
                2.3272683588860026e01,
                2.7223982220959247e01,
                -2.8240595076334000e02,
                3.6943078590316281e02,
                -1.3645364576977221e02,
                1.9058236733002300e01,
                2.1644988962398710e01,
                2.5012267757287322e01,
                -2.6180071928343307e02,
                3.4282650121799617e02,
                -1.2669036882336400e02,
                1.7955886187119816e01,
                2.3272683588634656e01,
                2.7223982224651898e01,
                -2.8240595077199526e02,
                3.6943078591032139e02,
                -1.3645364577174797e02,
                1.9058414960148450e01,
                2.1643466247439289e01,
                2.5016983354038196e01,
                -2.6179767020610126e02,
                3.4281320617581565e02,
                -1.2668337355331974e02,
                1.9023741366983238e01,
                2.5957710504548576e01,
                3.2054387652193789e01,
                -3.1898571318422574e02,
                4.1454655650462962e02,
                -1.5250373535684176e02,
                2.0050906563887416e01,
                2.4121527381838824e01,
                2.9665428981325245e01,
                -2.9589665055055406e02,
                3.8453661583827250e02,
                -1.4143340987287985e02,
                1.9023741366984520e01,
                2.5957710504498362e01,
                3.2054387653114766e01,
                -3.1898571318642672e02,
                4.1454655650647550e02,
                -1.5250373535735841e02,
                2.0051026978020587e01,
                2.4120379273875816e01,
                2.9669474257430963e01,
                -2.9589543070583102e02,
                3.8452729731205977e02,
                -1.4142824748467820e02,
                2.0092947487287756e01,
                2.8776895490568755e01,
                3.7339233558876920e01,
                -9.8781982607414882e00,
                7.0916635282296292e-01,
                -1.2340880155534291e-02,
                2.1044418341890132e01,
                2.6741847681518077e01,
                3.4861073630499796e01,
                -9.1700568642165461e00,
                6.5220324713443967e-01,
                -1.1045071585279443e-02,
                2.0092947487288011e01,
                2.8776895490557653e01,
                3.7339233559103448e01,
                -9.8781982608033179e00,
                7.0916635282857932e-01,
                -1.2340880155703077e-02,
                2.1044499630877905e01,
                2.6740987496092696e01,
                3.4864491165514394e01,
                -9.1707199731434574e00,
                6.5223741134844682e-01,
                -1.1045188698410773e-02,
                2.1162510215379026e01,
                3.1723491960797684e01,
                4.3084295875067085e01,
                -4.1033675985379521e00,
                -6.6095139594000130e-01,
                6.0977735530407223e-02,
                2.2038706806958309e01,
                2.9505670300337073e01,
                4.0630600131872811e01,
                -2.7905442844326718e00,
                -8.3885972791335117e-01,
                6.8309956404426039e-02,
                2.1162510215379076e01,
                3.1723491960795304e01,
                4.3084295875120795e01,
                -4.1033675985539224e00,
                -6.6095139593840913e-01,
                6.0977735530354210e-02,
                2.2038761643178379e01,
                2.9505029336592230e01,
                4.0633451796171073e01,
                -2.7913314472201640e00,
                -8.3878528163749511e-01,
                6.8307595298566767e-02,
                3.1719012432820758e01,
                6.7480322661109355e01,
                1.3318978565899991e02,
                -1.6791944323404795e01,
                -1.0181217992701848e00,
                1.2989592638281225e-01,
                3.2009499874031789e01,
                6.5013296175889408e01,
                1.3669799889514238e02,
                -1.7009031615065428e01,
                -1.0689880784706638e00,
                1.3388972346122466e-01,
                3.1719012432820758e01,
                6.7480322661109355e01,
                1.3318978565899991e02,
                -1.6791944323404795e01,
                -1.0181217992701848e00,
                1.2989592638281225e-01,
                3.2009500887769519e01,
                6.5013269472322307e01,
                1.3669829238273672e02,
                -1.7009116366540379e01,
                -1.0689798256828462e00,
                1.3388945486998777e-01,
                4.1931127118492086e01,
                1.1600186087954401e02,
                3.1751764022286790e02,
                -4.6438894455748802e01,
                -8.7599401950869438e-01,
                2.2297105562740663e-01,
                4.2002297497564768e01,
                1.1479764873768737e02,
                3.2393143797302810e02,
                -4.7847299173836262e01,
                -7.8150712905299369e-01,
                2.2131248436241077e-01,
                4.1931127118492086e01,
                1.1600186087954401e02,
                3.1751764022286790e02,
                -4.6438894455748802e01,
                -8.7599401950869438e-01,
                2.2297105562740663e-01,
                4.2002297514594851e01,
                1.1479764793294436e02,
                3.2393145467669495e02,
                -4.7847304068128608e01,
                -7.8150664807362491e-01,
                2.2131246858403722e-01,
                5.1984670105634827e01,
                1.7926303194781252e02,
                6.2846495111925287e02,
                -1.0034649475039414e02,
                2.4606292097951082e-01,
                3.3256752105517051e-01,
                5.2000554052128159e01,
                1.7883235795593501e02,
                6.3273302895025176e02,
                -1.0138733878813618e02,
                3.2804187851642969e-01,
                3.3055293107858102e-01,
                5.1984670105634827e01,
                1.7926303194781252e02,
                6.2846495111925287e02,
                -1.0034649475039414e02,
                2.4606292097951082e-01,
                3.3256752105517051e-01,
                5.2000554052402805e01,
                1.7883235793562420e02,
                6.3273302962903426e02,
                -1.0138733898825184e02,
                3.2804189825766372e-01,
                3.3055293042886030e-01,
                6.1996666427075382e01,
                2.5724136589119979e02,
                1.0913830717468406e03,
                -1.8317243758181812e02,
                2.5193786568880601e00,
                4.6277932792022042e-01,
                6.2000133522892554e01,
                2.5710536851489377e02,
                1.0934673032018356e03,
                -1.8370056934287794e02,
                2.5630609198690104e00,
                4.6162176037505448e-01,
                6.1996666427075382e01,
                2.5724136589119979e02,
                1.0913830717468406e03,
                -1.8317243758181812e02,
                2.5193786568880601e00,
                4.6277932792022042e-01,
                6.2000133522896938e01,
                2.5710536851442714e02,
                1.0934673032246803e03,
                -1.8370056934963364e02,
                2.5630609205366826e00,
                4.6162176035304603e-01,
                7.1999279107664492e01,
                3.4965254984584158e02,
                1.7356304176273381e03,
                -3.0063395678020430e02,
                6.2079056750108883e00,
                6.1505333334154833e-01,
                7.2000032172982571e01,
                3.4961232791697932e02,
                1.7365043785874466e03,
                -3.0086002522613632e02,
                6.2270725229979789e00,
                6.1452738833821030e-01,
                7.1999279107664492e01,
                3.4965254984584158e02,
                1.7356304176273381e03,
                -3.0063395678020430e02,
                6.2079056750108883e00,
                6.1505333334154833e-01,
                7.2000032172982642e01,
                3.4961232791696904e02,
                1.7365043785881401e03,
                -3.0086002522634379e02,
                6.2270725230187063e00,
                6.1452738833751985e-01,
                8.1999844359310714e01,
                4.5636323545227941e02,
                2.5918884526432239e03,
                -4.5885344883307727e02,
                1.1616256691917803e01,
                7.8948404417119522e-01,
                8.2000007751936337e01,
                4.5635184072744744e02,
                2.5922210189842476e03,
                -4.5894061525528980e02,
                1.1623761628208563e01,
                7.8927378661620728e-01,
                8.1999844359310714e01,
                4.5636323545227941e02,
                2.5918884526432239e03,
                -4.5885344883307727e02,
                1.1616256691917803e01,
                7.8948404417119522e-01,
                8.2000007751936337e01,
                4.5635184072744744e02,
                2.5922210189842476e03,
                -4.5894061525528980e02,
                1.1623761628208563e01,
                7.8927378661620728e-01,
            ],
            dtype=dtype,
            device=env.DEVICE,
        ).reshape(8, 174)  # 1392
        self.table_info_tensor = torch.tensor(
            [
                -2.1000000000000000e01,
                2.1000000000000000e01,
                1.0500000000000000e02,
                1.0000000000000000e00,
                1.0000000000000000e01,
                -1.0000000000000000e00,
            ],
            dtype=dtype,
            device="cpu",
        )
        self.em_x_tensor = torch.tensor(
            [
                9.3816147034272368e-01,
                -1.6703373029862567e-01,
                -4.4294526064601734e-02,
                -2.8798505489184573e-01,
                -1.6703373029862567e-01,
                9.2489218226366088e-01,
                -2.8928196536572048e-01,
                -4.7833509099876154e-01,
                -4.4294526064601734e-02,
                -2.8928196536572048e-01,
                5.7034320185695120e-01,
                1.8771147911830000e-01,
                -2.8798505489184573e-01,
                -4.7833509099876154e-01,
                1.8771147911830000e-01,
                4.0174654365823070e-01,
                8.4370316144902313e-01,
                -3.7813146789689916e-02,
                -3.6989397568296523e-01,
                -4.0554075086539937e-01,
                -3.7813146789689916e-02,
                6.5766402633747112e-01,
                -4.2312966361682885e-01,
                1.2685067374257861e-01,
                -3.6989397568296523e-01,
                -4.2312966361682885e-01,
                6.0171576901660107e-01,
                9.8283160997298613e-02,
                -4.0554075086539937e-01,
                1.2685067374257861e-01,
                9.8283160997298613e-02,
                2.1324148100625978e-01,
                9.7843596341516559e-01,
                -1.0492833888237871e-01,
                -1.0538688914576379e-01,
                -2.0453551592353389e-01,
                -1.0492833888237871e-01,
                7.7943976693565231e-01,
                -1.5898500035781410e-01,
                9.4834209331437741e-02,
                -1.0538688914576379e-01,
                -1.5898500035781410e-01,
                7.4778071691708869e-01,
                -6.1895255142095873e-01,
                -2.0453551592353389e-01,
                9.4834209331437741e-02,
                -6.1895255142095873e-01,
                6.0844713798743799e-01,
                1.0079020879244640e00,
                -2.3855984150631487e-01,
                -3.4608276043004524e-02,
                -4.7448768267289088e-01,
                -2.3855984150631487e-01,
                4.9732018171028253e-01,
                -3.1320787082485729e-01,
                -1.4528004145602180e-01,
                -3.4608276043004524e-02,
                -3.1320787082485729e-01,
                4.7696729363954582e-01,
                1.1723268074231248e-01,
                -4.7448768267289088e-01,
                -1.4528004145602180e-01,
                1.1723268074231248e-01,
                4.0511515406019899e-01,
            ],
            dtype=dtype,
            device=env.DEVICE,
        ).reshape(4, 16)  # 3072
        self.em_tensor = self.em_x_tensor.reshape(4, 4, 4)
        self.table_info_tensor.requires_grad = False
        self.table_tensor.requires_grad = False
        self.em_x_tensor.requires_grad = True
        self.em_tensor.requires_grad = True
        self.last_layer_size = 4
        self.nloc = 192
        self.nnei_i = 4
        self.nnei_j = 4

        self.expected_descriptor_tensor = torch.tensor(
            [
                1.4271973325754339e00,
                2.5214997685364109e00,
                3.1394341134078902e00,
                2.2727894815158436e00,
                1.9127738317829568e00,
                2.5288382955492263e00,
                3.1401587802428659e00,
                2.5252400661016079e00,
                9.4806287131835343e-01,
                2.3778589851963829e00,
                2.8273548699126683e00,
                1.9358633427396228e00,
                2.1586806210305824e00,
                2.6256636737020518e00,
                3.3955783231847523e00,
                2.7091329174140033e00,
            ],
            dtype=dtype,
            device=env.DEVICE,
        ).reshape(4, 4)
        self.expected_dy_dem_x = torch.tensor(
            [
                5.648489055364202,
                -0.8109841888364551,
                -0.24536867097411239,
                -1.1747441933374314,
                -0.8109841888364551,
                5.551778760144183,
                -1.177917429853053,
                -1.679018415609313,
                -0.24536867097411239,
                -1.177917429853053,
                3.066855971667982,
                1.0527786223200397,
                -1.1747441933374314,
                -1.679018415609313,
                1.0527786223200397,
                2.154128070312613,
                4.941344648306369,
                -0.2103366776681705,
                -1.367487115395829,
                -1.4564034402591373,
                -0.2103366776681705,
                3.608285604419272,
                -1.5040667548923485,
                0.7193959503618953,
                -1.367487115395829,
                -1.5040667548923485,
                3.254731998053089,
                0.5591533452715003,
                -1.4564034402591373,
                0.7193959503618953,
                0.5591533452715003,
                1.1890751500286143,
                5.930646543070966,
                -0.5504936332829675,
                -0.5526209441647617,
                -0.9416633243620358,
                -0.5504936332829675,
                4.459184963787069,
                -0.7803876483286257,
                0.5396747709549512,
                -0.5526209441647617,
                -0.7803876483286257,
                4.228294605542883,
                -2.4024371569601737,
                -0.9416633243620358,
                0.5396747709549512,
                -2.4024371569601737,
                3.2959904931493385,
                6.1224119486112745,
                -1.0448644883385283,
                -0.1928789305436163,
                -1.665305701289531,
                -1.0448644883385283,
                2.654867798405467,
                -1.2350917789985936,
                -0.7261316753808301,
                -0.1928789305436163,
                -1.2350917789985936,
                2.5453318630183253,
                0.6656811038445796,
                -1.665305701289531,
                -0.7261316753808301,
                0.6656811038445796,
                2.1713171576639834,
            ],
            dtype=dtype,
            device=env.DEVICE,
        ).reshape(4, 16)
        self.expected_dy_dem = torch.tensor(
            [
                -5.714759600210596,
                -11.831713987629353,
                -11.190122813510595,
                -12.370780785045307,
                -11.831713987629353,
                -5.794533094540567,
                -12.376066372858583,
                -13.075171860091254,
                -11.190122813510595,
                -12.376066372858583,
                -7.803917285588026,
                -9.877251464178656,
                -12.370780785045307,
                -13.075171860091254,
                -9.877251464178656,
                -8.705541422983027,
                -6.27628449412496,
                -11.154143613489751,
                -12.688057210800388,
                -12.817836346604595,
                -11.154143613489751,
                -7.330064381543641,
                -12.8806650535401,
                -10.220630213188823,
                -12.688057210800388,
                -12.8806650535401,
                -7.634743369735131,
                -10.382921713303563,
                -12.817836346604595,
                -10.220630213188823,
                -10.382921713303563,
                -9.734474315546018,
                -5.471377616034635,
                -11.517744760556846,
                -11.520149882632618,
                -12.009085059484455,
                -11.517744760556846,
                -6.648325517174191,
                -11.792420704526007,
                -10.402546165885312,
                -11.520149882632618,
                -11.792420704526007,
                -6.828388581370712,
                -13.58699904400806,
                -12.009085059484455,
                -10.402546165885312,
                -13.58699904400806,
                -7.598306221049991,
                -5.292546900481144,
                -12.161894405917199,
                -11.136299195212807,
                -13.061667970148363,
                -12.161894405917199,
                -8.194852173884833,
                -12.47192675731089,
                -11.724533065186144,
                -11.136299195212807,
                -12.47192675731089,
                -8.30347749201922,
                -10.275210657956574,
                -13.061667970148363,
                -11.724533065186144,
                -10.275210657956574,
                -8.687482898190318,
            ],
            dtype=dtype,
            device=env.DEVICE,
        ).reshape(4, 4, 4)

    def test_forward(self) -> None:
        # Call the forward function
        forward_result = torch.ops.deepmd.tabulate_fusion_se_t(
            self.table_tensor,
            self.table_info_tensor,
            self.em_x_tensor,
            self.em_tensor,
            self.last_layer_size,
        )

        descriptor_tensor = forward_result[0]

        # Check the shape
        self.assertEqual(descriptor_tensor.shape, self.expected_descriptor_tensor.shape)

        # Check the values
        torch.testing.assert_close(
            descriptor_tensor,
            self.expected_descriptor_tensor,
            atol=self.prec,
            rtol=self.prec,
        )

    def test_backward(self) -> None:
        # Call the forward function
        forward_result = torch.ops.deepmd.tabulate_fusion_se_t(
            self.table_tensor,
            self.table_info_tensor,
            self.em_x_tensor,
            self.em_tensor,
            self.last_layer_size,
        )

        descriptor_tensor = forward_result[0]

        # Check the forward
        torch.testing.assert_close(
            descriptor_tensor,
            self.expected_descriptor_tensor,
            atol=self.prec,
            rtol=self.prec,
        )

        # Create a loss and perform backward
        loss = descriptor_tensor.sum()
        loss.backward()

        # Check gradients
        self.assertIsNotNone(self.em_x_tensor.grad)
        self.assertIsNotNone(self.em_tensor.grad)

        # Check the shapes of the gradients
        self.assertEqual(self.em_x_tensor.grad.shape, self.expected_dy_dem_x.shape)
        self.assertEqual(self.em_tensor.grad.shape, self.expected_dy_dem.shape)

        # Check the values of the gradients
        torch.testing.assert_close(
            self.em_x_tensor.grad,
            self.expected_dy_dem_x,
            atol=self.prec,
            rtol=self.prec,
        )

        torch.testing.assert_close(
            self.em_tensor.grad,
            self.expected_dy_dem,
            atol=self.prec,
            rtol=self.prec,
        )


if __name__ == "__main__":
    unittest.main()
